{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2a8bbda-9a0c-4a99-971a-c311a8f3bbd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"]=\"0\"\n",
    "from peft import PeftModel, PeftConfig\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, TextStreamer\n",
    "import torch\n",
    "from threading import Thread\n",
    "import gradio as gr\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4c9da820-1c67-4cff-a9c1-7bb8201a183e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "719c963e5c18432582865343bbeeadac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/2 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PeftModelForCausalLM(\n",
       "  (base_model): LoraModel(\n",
       "    (model): MistralForCausalLM(\n",
       "      (model): MistralModel(\n",
       "        (embed_tokens): Embedding(32008, 4096)\n",
       "        (layers): ModuleList(\n",
       "          (0-31): 32 x MistralDecoderLayer(\n",
       "            (self_attn): MistralSdpaAttention(\n",
       "              (q_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=16, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=16, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (k_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=16, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=16, out_features=1024, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (v_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=1024, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=16, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=16, out_features=1024, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (o_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=16, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=16, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (rotary_emb): MistralRotaryEmbedding()\n",
       "            )\n",
       "            (mlp): MistralMLP(\n",
       "              (gate_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=16, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=16, out_features=14336, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (up_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=4096, out_features=14336, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=4096, out_features=16, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=16, out_features=14336, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (down_proj): lora.Linear4bit(\n",
       "                (base_layer): Linear4bit(in_features=14336, out_features=4096, bias=False)\n",
       "                (lora_dropout): ModuleDict(\n",
       "                  (default): Dropout(p=0.1, inplace=False)\n",
       "                )\n",
       "                (lora_A): ModuleDict(\n",
       "                  (default): Linear(in_features=14336, out_features=16, bias=False)\n",
       "                )\n",
       "                (lora_B): ModuleDict(\n",
       "                  (default): Linear(in_features=16, out_features=4096, bias=False)\n",
       "                )\n",
       "                (lora_embedding_A): ParameterDict()\n",
       "                (lora_embedding_B): ParameterDict()\n",
       "              )\n",
       "              (act_fn): SiLU()\n",
       "            )\n",
       "            (input_layernorm): MistralRMSNorm()\n",
       "            (post_attention_layernorm): MistralRMSNorm()\n",
       "          )\n",
       "        )\n",
       "        (norm): MistralRMSNorm()\n",
       "      )\n",
       "      (lm_head): Linear(in_features=4096, out_features=32008, bias=False)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "peft_model_id = \"smangrul/mistral-dpo\" #\"smangrul/mistral-sft\"\n",
    "config = PeftConfig.from_pretrained(peft_model_id)\n",
    "model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, load_in_4bit=True, device_map=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(peft_model_id)\n",
    "model.resize_token_embeddings(len(tokenizer), pad_to_multiple_of=8)\n",
    "model = PeftModel.from_pretrained(model, peft_model_id)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b45b93d1-e369-4f84-ae9a-d6f033e3f30f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class ChatCompletion:\n",
    "  def __init__(self, model, tokenizer, system_prompt=None):\n",
    "    self.model = model\n",
    "    self.tokenizer = tokenizer\n",
    "    self.streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True)\n",
    "    self.print_streamer = TextStreamer(self.tokenizer, skip_prompt=True)\n",
    "    # set the model in inference mode\n",
    "    self.model.eval()\n",
    "    self.system_prompt = system_prompt\n",
    "\n",
    "  def get_completion(self, prompt, system_prompt=None, message_history=None, max_new_tokens=512, temperature=0.0):\n",
    "    if temperature < 1e-2:\n",
    "      temperature = 1e-2\n",
    "    messages = []\n",
    "    if message_history is not None:\n",
    "      messages.extend(message_history)\n",
    "    elif system_prompt or self.system_prompt:\n",
    "      system_prompt = system_prompt or self.system_prompt\n",
    "      messages.append({\"role\": \"system\", \"content\":system_prompt})\n",
    "    messages.append({\"role\": \"user\", \"content\": prompt})\n",
    "    chat_prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "\n",
    "    inputs = self.tokenizer(chat_prompt, return_tensors=\"pt\", add_special_tokens=False)\n",
    "    # Run the generation in a separate thread, so that we can fetch the generated text in a non-blocking way.\n",
    "    generation_kwargs = dict(max_new_tokens=max_new_tokens,\n",
    "                             temperature=temperature,\n",
    "                             top_p=0.95,\n",
    "                             do_sample=True,\n",
    "                             repetition_penalty=1.1)\n",
    "    generated_text = self.model.generate(**inputs, streamer=self.print_streamer, **generation_kwargs)\n",
    "    return generated_text\n",
    "\n",
    "  def get_chat_completion(self, message, history):\n",
    "    messages = []\n",
    "    if self.system_prompt:\n",
    "      messages.append({\"role\": \"system\", \"content\":self.system_prompt})\n",
    "    for user_message, assistant_message in history:\n",
    "        messages.append({\"role\": \"user\", \"content\": user_message})\n",
    "        messages.append({\"role\": \"system\", \"content\": assistant_message})\n",
    "    messages.append({\"role\": \"user\", \"content\": message})\n",
    "    chat_prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "\n",
    "    inputs = self.tokenizer(chat_prompt, return_tensors=\"pt\")\n",
    "    # Run the generation in a separate thread, so that we can fetch the generated text in a non-blocking way.\n",
    "    generation_kwargs = dict(inputs,\n",
    "                             streamer=self.streamer,\n",
    "                             max_new_tokens=2048,\n",
    "                             temperature=0.2,\n",
    "                             top_p=0.95,\n",
    "                             do_sample=True,\n",
    "                             repetition_penalty=1.2,\n",
    "                             eos_token_id=tokenizer.eos_token_id\n",
    "                            )\n",
    "    thread = Thread(target=self.model.generate, kwargs=generation_kwargs)\n",
    "    thread.start()\n",
    "    generated_text = \"\"\n",
    "    for new_text in self.streamer:\n",
    "        generated_text += new_text\n",
    "        yield generated_text.replace(tokenizer.eos_token, \"\")\n",
    "    thread.join()\n",
    "    return generated_text.replace(tokenizer.eos_token, \"\")\n",
    "\n",
    "  def get_completion_without_streaming(self, prompt, system_prompt=None, message_history=None, max_new_tokens=512, temperature=0.0):\n",
    "    if temperature < 1e-2:\n",
    "      temperature = 1e-2\n",
    "    messages = []\n",
    "    if message_history is not None:\n",
    "      messages.extend(message_history)\n",
    "    elif system_prompt or self.system_prompt:\n",
    "      system_prompt = system_prompt or self.system_prompt\n",
    "      messages.append({\"role\": \"system\", \"content\":system_prompt})\n",
    "    messages.append({\"role\": \"user\", \"content\": prompt})\n",
    "    chat_prompt = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "\n",
    "    inputs = self.tokenizer(chat_prompt, return_tensors=\"pt\", add_special_tokens=False)\n",
    "    # Run the generation in a separate thread, so that we can fetch the generated text in a non-blocking way.\n",
    "    generation_kwargs = dict(max_new_tokens=max_new_tokens,\n",
    "                             temperature=temperature,\n",
    "                             top_p=0.95,\n",
    "                             do_sample=True,\n",
    "                             repetition_penalty=1.1)\n",
    "    outputs = self.model.generate(**inputs, **generation_kwargs)\n",
    "    generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)\n",
    "    return generated_text\n",
    "\n",
    "system_prompt = \"You are a helpful, honest and harmless chatbot named Rudra.\"\n",
    "text_generator = ChatCompletion(model, tokenizer, system_prompt=system_prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5c15652-9392-49ff-b2cf-3695a45bc29f",
   "metadata": {},
   "outputs": [],
   "source": [
    "gr.ChatInterface(text_generator.get_chat_completion).queue().launch(debug=True, share=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21f9c415-8059-454b-8a90-cb91803abd2f",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
